import asyncio
import logging
import os
import time
import uuid
import json
from typing import Dict, Any, Optional
from datetime import datetime

from three_d_effect_evaluation.model_evaluator import ModelEvaluator
from db_manager import DatabaseManager

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("model_processing_service")

class ModelProcessingService:
    """
    3D模型处理服务，负责在模型生成完成后进行自动评估和处理
    """
    
    def __init__(self):
        self.db_manager = DatabaseManager()
        self.model_evaluator = ModelEvaluator()
        self.processing_tasks = {}
        self.task_statuses = {}
    
    async def process_generated_model(self, job_id: str, user_id: int, prompt: str = None, reference_image: str = None) -> Dict[str, Any]:
        """
        处理生成完成的3D模型，执行评估流程
        
        Args:
            job_id: 任务ID
            user_id: 用户ID
            prompt: 提示词（文本生成时使用）
            reference_image: 参考图片（图片生成时使用）
            
        Returns:
            dict: 处理结果
        """
        try:
            self.task_statuses[job_id] = {"status": "PROCESSING", "step": "STARTING"}
            logger.info(f"开始处理模型: job_id={job_id}, user_id={user_id}")
            
            # 1. 检查任务状态
            self.task_statuses[job_id]["step"] = "CHECKING_STATUS"
            task_status = await self._check_task_status(job_id)
            if not task_status or task_status.get("Status") != "DONE":
                error_msg = f"任务未完成或状态异常: {task_status.get('Status', 'UNKNOWN')}"
                logger.error(error_msg)
                self.task_statuses[job_id] = {"status": "FAILED", "error": error_msg}
                return {"code": 0, "msg": error_msg, "data": {}}
            
            # 2. 获取模型URL
            self.task_statuses[job_id]["step"] = "EXTRACTING_MODEL_URL"
            model_url = await self._extract_model_url(task_status)
            if not model_url:
                error_msg = "无法从任务结果中提取模型URL"
                logger.error(error_msg)
                self.task_statuses[job_id] = {"status": "FAILED", "error": error_msg}
                return {"code": 0, "msg": error_msg, "data": {}}
            
            # 3. 执行模型评估
            self.task_statuses[job_id]["step"] = "EVALUATING_MODEL"
            evaluation_result = await self._evaluate_model(
                model_url=model_url,
                prompt=prompt,
                reference_image=reference_image
            )
            
            if not evaluation_result.get("success"):
                error_msg = f"模型评估失败: {evaluation_result.get('error', '未知错误')}"
                logger.error(error_msg)
                self.task_statuses[job_id] = {"status": "FAILED", "error": error_msg}
                return {"code": 0, "msg": error_msg, "data": {}}
            
            # 4. 保存评估报告
            self.task_statuses[job_id]["step"] = "SAVING_REPORT"
            report_data = {
                "title": self._generate_report_title(prompt, reference_image),
                "markdown_report": evaluation_result["report"]
            }
            
            save_result = await self._save_evaluation_report(
                model_3d_url=model_url,
                user_id=user_id,
                report_data=report_data,
                screenshot=evaluation_result.get("screenshot")
            )
            
            if not save_result:
                error_msg = "保存评估报告失败"
                logger.error(error_msg)
                self.task_statuses[job_id] = {"status": "FAILED", "error": error_msg}
                return {"code": 0, "msg": error_msg, "data": {}}
            
            # 5. 更新任务状态
            self.task_statuses[job_id]["step"] = "UPDATING_STATUS"
            await self._update_task_status(job_id, "COMPLETED")
            
            self.task_statuses[job_id] = {"status": "COMPLETED", "report_id": save_result["id"]}
            logger.info(f"模型处理完成: job_id={job_id}, report_id={save_result['id']}")
            
            return {
                "code": 1,
                "msg": "模型评估完成",
                "data": {
                    "job_id": job_id,
                    "model_url": model_url,
                    "report_id": save_result["id"],
                    "report_title": report_data["title"]
                }
            }
            
        except Exception as e:
            error_msg = f"处理模型时发生异常: {str(e)}"
            logger.error(error_msg, exc_info=True)
            self.task_statuses[job_id] = {"status": "FAILED", "error": str(e)}
            return {"code": 0, "msg": error_msg, "data": {}}
    
    async def _check_task_status(self, job_id: str) -> Optional[Dict[str, Any]]:
        """
        检查任务状态
        """
        try:
            # 从数据库中获取任务状态
            task_info = self.db_manager.get_task_by_id(job_id)
            if task_info and task_info.get('status') == 'DONE':
                # 构建符合处理逻辑的返回格式
                return {
                    'Status': 'DONE',
                    'ResultFile3Ds': [
                        {
                            'Type': 'model',
                            'Url': task_info.get('model_url', '')
                        }
                    ]
                }
            return {'Status': task_info.get('status', 'UNKNOWN')} if task_info else None
        except Exception as e:
            logger.error(f"检查任务状态失败: {str(e)}")
            return None
    
    async def _extract_model_url(self, task_status: Dict[str, Any]) -> Optional[str]:
        """
        从任务状态中提取3D模型URL
        """
        try:
            if 'ResultFile3Ds' in task_status:
                # 查找合适的3D模型文件
                for file_info in task_status['ResultFile3Ds']:
                    if file_info.get('Type') == 'model' and 'Url' in file_info:
                        return file_info['Url']
                    if 'Url' in file_info:
                        # 如果没有明确的model类型，返回第一个URL
                        return file_info['Url']
            return None
        except Exception as e:
            logger.error(f"提取模型URL失败: {str(e)}")
            return None
    
    async def _evaluate_model(self, model_url: str, prompt: str = None, reference_image: str = None) -> Dict[str, Any]:
        """
        评估3D模型
        """
        try:
            # 调用模型评估器进行评估
            result = await self.model_evaluator.evaluate_model(
                model_url=model_url,
                prompt=prompt,
                reference_image=reference_image
            )
            return result
        except Exception as e:
            logger.error(f"评估模型失败: {str(e)}")
            return {"success": False, "error": str(e)}
    
    async def _save_evaluation_report(self, model_3d_url: str, user_id: int, report_data: Dict[str, Any], screenshot: Optional[str] = None) -> Optional[Dict[str, Any]]:
        """
        保存评估报告到数据库
        """
        try:
            # 调用数据库管理器保存报告
            result = self.db_manager.insert_evaluation_report(
                model_3d_url=model_3d_url,
                user_id=user_id,
                title=report_data["title"],
                markdown_report=report_data["markdown_report"],
                screenshot=screenshot
            )
            return result
        except Exception as e:
            logger.error(f"保存评估报告失败: {str(e)}")
            return None
    
    async def _update_task_status(self, job_id: str, status: str) -> bool:
        """
        更新任务状态
        """
        try:
            # 更新数据库中的任务状态
            return self.db_manager.update_generation_status(task_id=job_id, status=status)
        except Exception as e:
            logger.error(f"更新任务状态失败: {str(e)}")
            return False
    
    def _generate_report_title(self, prompt: str = None, reference_image: str = None) -> str:
        """
        生成报告标题
        """
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        if prompt:
            # 截取提示词的前30个字符作为标题的一部分
            short_prompt = prompt[:30] + "..." if len(prompt) > 30 else prompt
            return f"{short_prompt}_评估报告_{timestamp}"
        elif reference_image:
            return f"基于图片的3D模型评估报告_{timestamp}"
        else:
            return f"3D模型评估报告_{timestamp}"
    
    def get_processing_status(self, job_id: str) -> Dict[str, Any]:
        """
        获取处理任务的状态
        
        Args:
            job_id: 任务ID
            
        Returns:
            dict: 任务状态信息
        """
        if job_id not in self.task_statuses:
            return {"status": "NOT_FOUND", "message": "任务未开始处理"}
        return self.task_statuses.get(job_id, {})
    
    async def process_model_with_sse(self, job_id: str, user_id: int, prompt: str = None, reference_image: str = None):
        """
        使用SSE方式处理模型并返回结果
        
        Args:
            job_id: 任务ID
            user_id: 用户ID
            prompt: 提示词
            reference_image: 参考图片
            
        Yields:
            str: SSE格式的事件数据
        """
        # 启动处理任务
        self.processing_tasks[job_id] = asyncio.create_task(
            self.process_generated_model(job_id, user_id, prompt, reference_image)
        )
        
        # 定期发送状态更新
        while job_id in self.task_statuses:
            status = self.task_statuses[job_id]
            yield {
                "event": "status",
                "data": status
            }
            
            # 如果处理完成或失败，跳出循环
            if status["status"] in ["COMPLETED", "FAILED"]:
                break
            
            # 每2秒发送一次更新
            await asyncio.sleep(2)
        
        # 等待任务完成并获取最终结果
        if job_id in self.processing_tasks:
            try:
                result = await self.processing_tasks[job_id]
                # 发送最终结果事件
                yield {
                    "event": "result",
                    "data": result
                }
            except Exception as e:
                yield {
                    "event": "error",
                    "data": {"code": 0, "msg": str(e), "data": {}}
                }
            finally:
                # 清理任务
                if job_id in self.processing_tasks:
                    del self.processing_tasks[job_id]

# 创建处理服务实例
model_processing_service = ModelProcessingService()